package services

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
)

// BatchOCRService 批量OCR服务
type BatchOCRService struct {
	siliconFlowService *SiliconFlowService
}

// NewBatchOCRService 创建批量OCR服务
func NewBatchOCRService(siliconFlowService *SiliconFlowService) *BatchOCRService {
	return &BatchOCRService{
		siliconFlowService: siliconFlowService,
	}
}

// FileOCRResult 单个文件的OCR结果
type FileOCRResult struct {
	FilePath string                 `json:"filePath"`
	FileName string                 `json:"fileName"`
	Success bool                    `json:"success"`
	Text    string                  `json:"text"`
	Data    map[string]interface{}  `json:"data"`
	Error   string                  `json:"error,omitempty"`
}

// BatchOCRResult 批量OCR结果
type BatchOCRResult struct {
	Total    int             `json:"total"`
	Success  int              `json:"success"`
	Failed   int              `json:"failed"`
	Results  []FileOCRResult  `json:"results"`
	MergedData map[string]interface{} `json:"mergedData"` // 合并所有OCR结果
}

// ScanImageFiles 扫描文件夹中的图片文件
func (s *BatchOCRService) ScanImageFiles(folderPath string) ([]string, error) {
	var imageFiles []string
	imageExts := map[string]bool{
		".jpg":  true,
		".jpeg": true,
		".png":  true,
		".bmp":  true,
		".gif":  true,
		".webp": true,
		".pdf":  true, // 暂时也支持PDF，但需要转换为图片
	}

	err := filepath.Walk(folderPath, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		if !info.IsDir() {
			ext := strings.ToLower(filepath.Ext(path))
			if imageExts[ext] {
				imageFiles = append(imageFiles, path)
			}
		}

		return nil
	})

	return imageFiles, err
}

// ReadImageAsBase64 读取图片文件并转换为base64
func (s *BatchOCRService) ReadImageAsBase64(filePath string) (string, error) {
	data, err := os.ReadFile(filePath)
	if err != nil {
		return "", err
	}

	base64Str := base64.StdEncoding.EncodeToString(data)
	return base64Str, nil
}

// RecognizeFolder 批量识别文件夹中的图片
func (s *BatchOCRService) RecognizeFolder(folderPath string, modelID string, onProgress func(current, total int, fileName string)) (*BatchOCRResult, error) {
	// 扫描图片文件
	imageFiles, err := s.ScanImageFiles(folderPath)
	if err != nil {
		return nil, fmt.Errorf("扫描文件夹失败: %v", err)
	}

	if len(imageFiles) == 0 {
		return &BatchOCRResult{
			Total:   0,
			Success: 0,
			Failed:  0,
			Results: []FileOCRResult{},
			MergedData: make(map[string]interface{}),
		}, nil
	}

	result := &BatchOCRResult{
		Total:     len(imageFiles),
		Success:   0,
		Failed:    0,
		Results:   make([]FileOCRResult, 0, len(imageFiles)),
		MergedData: make(map[string]interface{}),
	}

	// 逐个识别
	for i, filePath := range imageFiles {
		fileName := filepath.Base(filePath)
		
		if onProgress != nil {
			onProgress(i+1, len(imageFiles), fileName)
		}

		fileResult := FileOCRResult{
			FilePath: filePath,
			FileName: fileName,
			Success:  false,
		}

		// 读取图片为base64
		imageBase64, err := s.ReadImageAsBase64(filePath)
		if err != nil {
			fileResult.Error = fmt.Sprintf("读取文件失败: %v", err)
			result.Results = append(result.Results, fileResult)
			result.Failed++
			continue
		}

		// OCR识别
		text, err := s.siliconFlowService.RecognizeImage(modelID, imageBase64)
		if err != nil {
			fileResult.Error = fmt.Sprintf("OCR识别失败: %v", err)
			result.Results = append(result.Results, fileResult)
			result.Failed++
			continue
		}

		fileResult.Success = true
		fileResult.Text = text

		// 尝试解析为JSON
		var data map[string]interface{}
		if err := json.Unmarshal([]byte(text), &data); err == nil {
			fileResult.Data = data
			// 合并到总数据中
			for k, v := range data {
				// 如果键已存在，尝试合并（这里简单处理，使用最后一个值）
				result.MergedData[k] = v
			}
		} else {
			// 如果不是JSON，尝试从文本中提取结构化数据
			fileResult.Data = s.extractDataFromText(text)
			for k, v := range fileResult.Data {
				result.MergedData[k] = v
			}
		}

		result.Results = append(result.Results, fileResult)
		result.Success++
	}

	return result, nil
}

// extractDataFromText 从文本中提取结构化数据
func (s *BatchOCRService) extractDataFromText(text string) map[string]interface{} {
	data := make(map[string]interface{})
	
	// 简单的文本解析：查找 "字段名：值" 或 "字段名:值" 格式
	lines := strings.Split(text, "\n")
	for _, line := range lines {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}

		// 尝试匹配 "字段名：值" 格式
		parts := strings.SplitN(line, "：", 2)
		if len(parts) != 2 {
			parts = strings.SplitN(line, ":", 2)
		}

		if len(parts) == 2 {
			key := strings.TrimSpace(parts[0])
			value := strings.TrimSpace(parts[1])
			if key != "" && value != "" {
				data[key] = value
			}
		}
	}

	return data
}
